from flask_restx import Resource, Namespace
from flask import request
import os
from models.database import get_db_connection
from utils.auth import token_required
from datetime import datetime
import uuid
from concurrent.futures import ThreadPoolExecutor
from services.baidu_service import BaiduAIService  # 导入百度服务
import time

ns = Namespace('batch', description='批量生成相关接口')

# 创建一个全局的 BaiduAIService 实例
baidu_service = BaiduAIService()

def generate_single_emoji(params):
    """单个表情包生成函数"""
    try:
        text = params.get('text')
        style = params.get('style', 'cartoon')
        
        # 调用百度AI生成表情包
        result = baidu_service.generate_emoji(
            text=text,
            style=style
        )
        
        # 检查错误码
        if isinstance(result, dict) and result.get('error_code') == 18:
            # QPS 限制，等待一下再重试
            time.sleep(1)  # 等待1秒
            result = baidu_service.generate_emoji(
                text=text,
                style=style
            )
        
        if result.get('status') == 'success' or (isinstance(result, dict) and 'image_url' in result):
            return {
                'status': 'success',
                'data': {
                    'text': text,
                    'image_url': result.get('image_url', '')
                }
            }
        else:
            print(f"Generation failed for text: {text}")
            print("API response:", result)
            return {
                'status': 'error',
                'message': '生成失败'
            }
            
    except Exception as e:
        print(f"Error generating emoji for text '{text}': {str(e)}")
        return {
            'status': 'error',
            'message': str(e)
        }

@ns.route('/generate')
class BatchGenerator(Resource):
    @token_required
    def post(self):
        """批量生成表情包"""
        try:
            data = request.get_json()
            texts = data.get('texts', [])
            style = data.get('style', 'cartoon')
            user_id = request.user.get('user_id')

            if not texts:
                return {'status': 'error', 'message': '请提供要生成的文本列表'}, 400

            if len(texts) > 5:  # 限制最大生成数量
                return {'status': 'error', 'message': '一次最多生成5个表情包'}, 400

            results = []
            conn = get_db_connection()
            cursor = conn.cursor()

            try:
                # 使用线程池并发生成，限制为2个并发
                with ThreadPoolExecutor(max_workers=2) as executor:  # 修改这里
                    futures = []
                    for text in texts:
                        if not text.strip():  # 跳过空文本
                            continue
                            
                        # 为每个生成任务创建一个唯一ID
                        history_id = str(uuid.uuid4())
                        
                        # 准备生成参数
                        params = {
                            'text': text,
                            'style': style,
                            'history_id': history_id,
                            'user_id': user_id
                        }
                        
                        # 提交生成任务
                        future = executor.submit(generate_single_emoji, params)
                        futures.append((future, history_id))
                        time.sleep(0.6)  # 增加间隔到0.6秒，确保QPS不超过2
                    
                    # 收集所有生成结果
                    successful_results = []  # 记录成功的结果
                    for future, history_id in futures:
                        try:
                            result = future.result()
                            if result.get('status') == 'success':
                                # 记录到历史表
                                cursor.execute("""
                                    INSERT INTO history (
                                        history_id, user_id, prompt_text, 
                                        style, output_url, created_at
                                    ) VALUES (%s, %s, %s, %s, %s, %s)
                                """, (
                                    history_id,
                                    user_id,
                                    result['data']['text'],
                                    style,
                                    result['data']['image_url'],
                                    datetime.now()
                                ))
                                
                                successful_results.append({
                                    'history_id': history_id,
                                    'text': result['data']['text'],
                                    'output_url': result['data']['image_url']
                                })
                            
                        except Exception as e:
                            print(f"Error processing result: {str(e)}")
                            continue

                conn.commit()
                
                # 只要有成功的结果就返回
                if successful_results:
                    return {
                        'status': 'success',
                        'data': {
                            'total': len(successful_results),
                            'results': successful_results
                        }
                    }
                else:
                    return {
                        'status': 'error',
                        'message': '生成失败，请重试'
                    }, 500

            except Exception as e:
                conn.rollback()
                print(f"Batch generation error: {str(e)}")
                return {
                    'status': 'error',
                    'message': '批量生成失败'
                }, 500
            
            finally:
                cursor.close()
                conn.close()

        except Exception as e:
            print(f"Batch API error: {str(e)}")
            return {
                'status': 'error',
                'message': '服务器错误'
            }, 500